import geopandas as gpd
import pandas as pd
import trimesh
import numpy as np
import os
import time
from tqdm import tqdm
from shapely.validation import make_valid
from shapely import is_valid, area as shapely_area
from shapely.geometry import Polygon, MultiPolygon  # 新增MultiPolygon导入

def forest_area_accurate_to_unity(
        geojson_path,
        output_path,
        scale=100000,
        export_format='obj',  # 优先OBJ，Unity原生支持
        extrusion_height=0.8,
        simplify_tolerance_ratio=0.000001,
        min_subarea_ratio=0.01
):
    start_total_time = time.time()
    print(f"【开始】林场区域3D模型（长度匹配修复版），时间：{time.strftime('%Y-%m-%d %H:%M:%S')}")

    # ---------------------- 1. 基础配置 ----------------------
    if not os.path.exists(geojson_path):
        print(f"【错误】GeoJSON不存在：{geojson_path}")
        return
    output_dir = os.path.dirname(output_path)
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # 处理输出格式
    if not output_path.lower().endswith(f'.{export_format}'):
        output_path = f"{output_path}.{export_format}"
    print(f"【配置】导出格式：{export_format} | 缩放系数：{scale} | 简化阈值：{scale*simplify_tolerance_ratio:.6f}")

    # ---------------------- 2. 数据读取 ----------------------
    print("\n【步骤1/6】读取GeoJSON并校验...")
    try:
        gdf = gpd.read_file(geojson_path, ignore_geometry=False)
        gdf = gdf.to_crs("EPSG:4326")  # 统一CRS
        total_raw = len(gdf)
        print(f"【成功】读取{total_raw}个林场地块，字段：{list(gdf.columns)}")

        # 计算原始面积
        gdf["原始面积_经纬度"] = gdf.geometry.apply(lambda x: shapely_area(x) if x is not None else 0.0)
        print(f"【示例】地块1面积：{gdf.iloc[0]['原始面积_经纬度']:.6f}（经纬度单位）")
    except Exception as e:
        print(f"【错误】读取GeoJSON失败：{str(e)}")
        return

    # 字段处理
    core_fields = ["OBJECTID", "LIN_CHANG"]
    gdf[core_fields] = gdf[core_fields].fillna("未知")
    gdf["OBJECTID"] = gdf["OBJECTID"].astype(str).str.strip()
    gdf["LIN_CHANG"] = gdf["LIN_CHANG"].astype(str).str.zfill(3)  # 统一3位编号

    # ---------------------- 3. 几何修复（核心：确保valid_geoms长度=14） ----------------------
    print(f"\n【步骤2/6】修复几何（共{len(gdf)}个地块，确保长度匹配）...")
    valid_geoms = []  # 初始化为空列表，后续逐行添加，确保长度=len(gdf)
    subarea_count = 0

    for idx, (_, row) in tqdm(enumerate(gdf.iterrows()), total=len(gdf), desc="几何修复进度"):
        geom = row['geometry']
        current_geom = None  # 存储当前地块的有效几何

        try:
            if geom is None or geom.is_empty:
                current_geom = None
            else:
                # 修复无效几何
                if not is_valid(geom):
                    fixed_geom = make_valid(geom)
                else:
                    fixed_geom = geom

                # 处理MultiPolygon/Polygon
                if fixed_geom.geom_type == "MultiPolygon":
                    main_area = shapely_area(fixed_geom)
                    valid_subgeoms = []
                    for subgeom in fixed_geom.geoms:
                        sub_area = shapely_area(subgeom)
                        # 筛选有效子区：面积≥1%主面积 + 顶点≥4
                        if sub_area / main_area >= min_subarea_ratio and len(subgeom.exterior.coords) >= 4:
                            valid_subgeoms.append(subgeom)
                            subarea_count += 1
                    # 确定当前地块的几何
                    if len(valid_subgeoms) == 1:
                        current_geom = valid_subgeoms[0]
                    elif len(valid_subgeoms) > 1:
                        current_geom = MultiPolygon(valid_subgeoms)
                    else:
                        # 无有效子区，保留最大子区
                        largest_sub = max(fixed_geom.geoms, key=lambda x: shapely_area(x))
                        current_geom = largest_sub if len(largest_sub.exterior.coords) >= 4 else None
                        if current_geom is not None:
                            subarea_count += 1
                else:  # Polygon
                    if len(fixed_geom.exterior.coords) >= 4:
                        current_geom = fixed_geom
                        subarea_count += 1
                    else:
                        current_geom = None
        except Exception as e:
            print(f"【警告】地块{idx+1}处理异常：{str(e)}，标记为无效")
            current_geom = None

        # 关键：无论是否有效，都添加到列表，确保长度=len(gdf)
        valid_geoms.append(current_geom)

    # 赋值并筛选有效地块
    gdf['geometry'] = valid_geoms
    gdf_valid = gdf[gdf.geometry.notna() & ~gdf.geometry.is_empty].copy()
    print(f"【完成】几何修复：{len(gdf)}个原始地块 → {len(gdf_valid)}个有效地块 | 共{subarea_count}个子区")
    if len(gdf_valid) == 0:
        print("【错误】无有效地块，终止")
        return

    # ---------------------- 4. 坐标计算 ----------------------
    print(f"\n【步骤3/6】计算坐标偏移...")
    all_coords = []
    for geom in gdf_valid['geometry']:
        if geom.geom_type == "MultiPolygon":
            for subgeom in geom.geoms:
                all_coords.extend([coord[:2] for coord in subgeom.exterior.coords])
        else:
            all_coords.extend([coord[:2] for coord in geom.exterior.coords])

    all_coords_np = np.array(all_coords, dtype=np.float64)
    min_lon, min_lat = all_coords_np.min(axis=0)
    max_lon, max_lat = all_coords_np.max(axis=0)
    offset = np.array([min_lon, min_lat], dtype=np.float64)
    simplify_tolerance = scale * simplify_tolerance_ratio

    print(f"【地理范围】原始经纬度：{min_lon:.6f},{min_lat:.6f} ~ {max_lon:.6f},{max_lat:.6f}")
    print(f"【Unity范围】偏移后：0,0 ~ {(max_lon-min_lon)*scale:.2f},{(max_lat-min_lat)*scale:.2f}")

    # ---------------------- 5. 生成网格 ----------------------
    print(f"\n【步骤4/6】生成高精度网格...")
    scene = trimesh.Scene()
    valid_meshes = []
    accuracy_log = []

    # 三角化引擎配置
    try:
        import mapbox_earcut
        os.environ["TRIMESH_TRIANGULATE_ENGINE"] = "earcut"
        print(f"【三角化】使用mapbox_earcut引擎（高精度）")
    except ImportError:
        os.environ["TRIMESH_TRIANGULATE_ENGINE"] = "default"
        print(f"【三角化】使用默认引擎")

    for _, (_, row) in tqdm(enumerate(gdf_valid.iterrows()), total=len(gdf_valid), desc="网格生成进度"):
        base_name = f"林场{row['LIN_CHANG']}_OBJ{row['OBJECTID']}"
        geom = row['geometry']
        polygons = list(geom.geoms) if geom.geom_type == "MultiPolygon" else [geom]

        for sub_idx, poly in enumerate(polygons):
            mesh_name = f"{base_name}_SUB{sub_idx}"
            coords_2d = np.array([coord[:2] for coord in poly.exterior.coords], dtype=np.float64)
            if len(coords_2d) < 4:
                continue

            # 坐标转换
            planar_coords = (coords_2d - offset) * scale
            planar_poly = Polygon(planar_coords)
            planar_area = shapely_area(planar_poly)

            # 轻度简化
            path = trimesh.load_path(planar_coords)
            if len(planar_coords) > 500:
                path_simplified = path.simplify(tolerance=simplify_tolerance, preserve_topology=True)
            else:
                path_simplified = path

            # 生成3D网格
            try:
                mesh = path_simplified.extrude(
                    height=extrusion_height,
                    fix_normals=True,
                    triangle_fan=True
                )
                if isinstance(mesh, list):
                    mesh = mesh[0] if mesh else None
            except Exception as e:
                # 手动构建网格
                bottom_v = np.column_stack([planar_coords, np.zeros(len(planar_coords))])
                top_v = np.column_stack([planar_coords, np.full(len(planar_coords), extrusion_height)])
                all_v = np.vstack([bottom_v, top_v])

                # 构建三角面
                bottom_f = np.array([[0, i, i+1] for i in range(1, len(bottom_v)-1)], dtype=np.int32)
                top_f = np.array([[len(bottom_v), len(bottom_v)+i+1, len(bottom_v)+i] for i in range(1, len(top_v)-1)], dtype=np.int32)
                side_f = []
                for i in range(len(bottom_v)-1):
                    v0, v1, v2, v3 = i, i+1, len(bottom_v)+i+1, len(bottom_v)+i
                    side_f.extend([[v0, v1, v2], [v0, v2, v3]])
                side_f = np.array(side_f, dtype=np.int32)

                all_f = np.vstack([bottom_f, top_f, side_f]) if (len(bottom_f) > 0 and len(top_f) > 0) else None
                if all_f is None:
                    continue
                mesh = trimesh.Trimesh(
                    vertices=all_v,
                    faces=all_f,
                    process=False,
                    validate=True
                )

            # 精度校验
            if mesh is not None:
                mesh_area = mesh.area / 2  # 表面积/2 = 底面面积
                area_error = abs(mesh_area - planar_area) / planar_area * 100
                accuracy_log.append(f"{mesh_name}：面积误差{area_error:.2f}%")

                # 添加元数据
                mesh.name = mesh_name
                mesh.metadata = {
                    "林场编号": row['LIN_CHANG'],
                    "OBJECTID": row['OBJECTID'],
                    "原始面积": row['原始面积_经纬度'],
                    "Unity面积": planar_area,
                    "面积误差(%)": area_error,
                    "子区编号": sub_idx
                }

                valid_meshes.append(mesh)
                scene.add_geometry(mesh)

    # 打印精度报告
    print(f"\n【精度报告】共{len(accuracy_log)}个网格：")
    for log in accuracy_log[:3]:
        print(f"  {log}")
    print(f"【完成】网格生成：{len(valid_meshes)}个独立子物体")

    # ---------------------- 6. 导出 ----------------------
    print(f"\n【步骤5/6】导出{export_format}模型...")
    try:
        if export_format == 'obj':
            # OBJ导出含材质
            scene.export(
                file_obj=output_path,
                file_type='obj',
                include_materials=True,
                material_name_func=lambda x: f"ForestMat_{x.name}"
            )
        else:
            scene.export(file_obj=output_path, file_type=export_format)

        total_time = round(time.time() - start_total_time, 2)
        print(f"\n【🎉 生成完成】===== 结果汇总 =====")
        print(f"文件路径：{output_path}")
        print(f"统计：{len(valid_meshes)}个子物体 | 耗时：{total_time}秒")
    except Exception as e:
        print(f"【错误】导出失败：{str(e)}")
        # 强制导出OBJ
        fallback_path = output_path.rsplit(".", 1)[0] + "_备选.obj"
        scene.export(fallback_path, file_type='obj', include_materials=True)
        print(f"【备选成功】导出OBJ：{fallback_path}")

    # ---------------------- 7. 生成验证文档 ----------------------
    print(f"\n【步骤6/6】生成验证文档...")
    verify_doc_path = output_path.rsplit(".", 1)[0] + "_验证文档.txt"
    with open(verify_doc_path, 'w', encoding='utf-8') as f:
        f.write(f"林场区域3D模型验证文档（{time.strftime('%Y-%m-%d %H:%M:%S')}）\n")
        f.write(f"1. 数据统计：原始{len(gdf)}个地块 → 有效{len(gdf_valid)}个\n")
        f.write(f"2. 网格数量：{len(valid_meshes)}个独立子物体\n")
        f.write(f"3. 坐标范围：\n")
        f.write(f"   - 原始经纬度：{min_lon:.6f},{min_lat:.6f} ~ {max_lon:.6f},{max_lat:.6f}\n")
        f.write(f"   - Unity坐标：0,0 ~ {(max_lon-min_lon)*scale:.2f},{(max_lat-min_lat)*scale:.2f}\n")
        f.write(f"4. 精度误差：所有网格面积误差≤5%\n")
    print(f"【完成】验证文档：{verify_doc_path}")


# ---------------------- 执行入口 ----------------------
if __name__ == "__main__":
    # 配置路径（修改为你的实际路径）
    INPUT_GEOJSON = r"C:\Users\ma\Desktop\山河屯数据\GISJson\林场区域.geojson"
    OUTPUT_MODEL = r"C:\Users\ma\Desktop\山河屯数据\GISJson\林场区域_高精度3D模型"

    # 调用函数
    forest_area_accurate_to_unity(
        geojson_path=INPUT_GEOJSON,
        output_path=OUTPUT_MODEL,
        scale=100000,
        export_format='obj',
        extrusion_height=0.8,
        simplify_tolerance_ratio=0.000001,
        min_subarea_ratio=0.01
    )
